#! /bin/bash

# Runs the "345M" parameter model
# 每个节点几个卡
NGPUS_PER_NODE=8
# 每个节点多少个流水线阶段
STAGES_PER_NODE=4
# 根据卡和流水线阶段算出来每个阶段多少个卡
GPUS_PER_SATGE=$(($NGPUS_PER_NODE/$STAGES_PER_NODE))
# Change for multinode config
# 主节点ip，多机跑需要设置
MASTER_ADDR=172.16.1.49 
MASTER_PORT=6001
# 多机设置节点数
NNODES=1
# 设置当前节点rank
NODE_RANK=0
# WORLD_SIZE=$(($STAGES_PER_NODE*$NNODES))
# add for test git in 12.10
DISTRIBUTED_ARGS="--nproc_per_node $STAGES_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT"

python -m torch.distributed.launch $DISTRIBUTED_ARGS \
       train_gpt.py \
       --tensor-model-parallel-size 1 \
       --pipeline-model-parallel-size $STAGES_PER_NODE \
       --epoch 1 \
       --gpus_per_stage $GPUS_PER_SATGE \
       --nstages_per_node $STAGES_PER_NODE \
       --micro-batch-size 16 \
       --micro-batch 16 \
       --distributed-backend nccl \
       --tensor-length 128


